# %load AI/code/trainval/class_train.py
from AI.utils.train_utils import get_accuracy, save_checkpoint, logger, adjust_learning_rate

def main():
    train_losses, val_losses, val_accuracyes = main_worker(15, gpu_flag=True)
    return train_losses, val_losses, val_accuracyes

def main_worker(num_epoch, gpu_flag=False):
    train_log = logger('runs//train')
    val_log = logger('runs//val')
    
    if args.resume:
        print("=> loading checkpoint '{}'".format(args.resume))
        if args.gpu is None:
            checkpoint = torch.load(args.resume)
        else:
            # Map model to be loaded to specified single gpu.
            loc = 'cuda:{}'.format(args.gpu)
            checkpoint = torch.load(args.resume, map_location=loc)
        args.start_epoch = checkpoint['epoch']
        best_acc1 = checkpoint['best_acc1']    
        model.load_state_dict(checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        print("=> loaded checkpoint '{}' (epoch {})"
              .format(args.resume, checkpoint['epoch']))
        
    train_losses, val_losses = [], []
    val_accuracyes = []
    best_acc1 = 0
    tmp_time=start_time = time.time()
    # 开始训练
    print('Start training!')
    for epoch in range(args.start_epoch, args.epochs):
        adjust_learning_rate(optimizer, epoch, args.lr, args.setlr)
        
        train_out = train_model(epoch, train_loader, model, criterion, optimizer, gpu_flag)
        val_out = val_model(epoch, val_loader, model, criterion, optimizer, gpu_flag)
        train_log.add_train_data(train_out, epoch)
        val_log.add_train_data(val_out, epoch)
        is_best = val_out[1] > best_acc1
        best_acc1 = max(val_out[1], best_acc1)
        save_checkpoint({
                'epoch': epoch + 1,
                'arch': model,
                'state_dict': model.state_dict(),
                'best_acc1' : best_acc1,
                'acc1': val_out[1],
                'acc5': val_out[2],
                'optimizer' : optimizer.state_dict(),
            },is_best, 'resnet50')

        train_losses.append(train_out[0])
        val_losses.append(val_out[0])
        val_accuracyes.append(val_out[1])
        print("Epoch: {epoch} Loss_avg：{loss:6.4f} \t Acc@1_avg {acc1:6.2f} \t Acc@5_avg {acc5:6.2f} \t \
Time@ {time:6.2f}\t TotalTime@{ttime:6.2f}"
                                                          .format(epoch=epoch,
                                                                  loss=train_out[0],
                                                                  acc1=val_out[1],
                                                                  acc5=val_out[2],
                                                                  time=time.time()-tmp_time,
                                                                  ttime=time.time()-start_time))
        tmp_time = time.time()
    print('Finished training!')
    train_log.close()
    val_log.close()
    
    return train_losses, val_losses, val_accuracyes
    
def train_model(epoch, data_loader, model, cirterion, optimizer, gpu_flag=False):
    '''
    模型的训练函数
    '''
    if gpu_flag:
        model.cuda().train()
    else:
        model.cpu().train()
        
    loss_avg = 0.0
    acc1_avg=0.0
    acc5_avg=0.0
    for i, data in enumerate(data_loader):
        images, labels = data
        if gpu_flag:
            images = Variable(images).cuda()
            labels = Variable(labels).cuda()
        else:
            images = Variable(images)
            labels = Variable(labels)
        optimizer.zero_grad()
        log_ps = model(images)
        loss = criterion(log_ps, labels)
        loss.backward()
        optimizer.step()
        acc1,  acc5 = get_accuracy(log_ps, labels, topk=(1, 5))
        acc1_avg += acc1.item()
        acc5_avg += acc5.item()
        loss_avg += loss.item()
        
        if i%500 == 499:
            print("Train Epoch: {epoch} Loss_avg：{loss:6.4f} \t Acc@1_avg {acc1:6.2f} \t Acc@5_avg {acc5:6.2f}".format(epoch=epoch,
                                                                  loss=loss_avg/(i+1),
                                                                  acc1=acc1_avg/(i+1),
                                                                  acc5=acc5_avg/(i+1) ))
    
    return loss_avg/len(data_loader),acc1_avg/len(data_loader),acc5_avg/len(data_loader)
        
    
def val_model(epoch, data_loader, model, cirterion, optimizer, gpu_flag=False):
    '''
    验证函数
    '''
    if gpu_flag:
        model.cuda().eval()
    else:
        model.cpu().eval()
    
    loss_avg = 0.0
    acc1_avg=0.0
    acc5_avg=0.0
    for i, data in enumerate(data_loader):
        images, labels = data
        if gpu_flag:
            images = Variable(images).cuda()
            labels = Variable(labels).cuda()
        else:
            images = Variable(images)
            labels = Variable(labels)
        log_ps = model(images)
        loss = criterion(log_ps, labels)
        acc1,  acc5 = get_accuracy(log_ps, labels, topk=(1, 5))
        acc1_avg += acc1.item()
        acc5_avg += acc5.item()
        loss_avg += loss.item()
        
        if i%500 == 499:
            print("Val   Epoch: {epoch} Loss_avg：{loss:6.4f} \t Acc@1_avg {acc1:6.2f} \t Acc@5_avg {acc5:6.2f}".format(epoch=epoch,
                                                                  loss=loss_avg/(i+1),
                                                                  acc1=acc1_avg/(i+1),
                                                                  acc5=acc5_avg/(i+1) ))
    
    return loss_avg/len(data_loader),acc1_avg/len(data_loader),acc5_avg/len(data_loader)
